import json
import re
import uuid

from flask import g

from myapp.app import app
from myapp.third.k8s import py_k8s
from myapp.utils.device_resource import DeviceResource

conf = app.config


def make_container_template(pipeline, task):
    # task = all_tasks[task_name]
    ops_args = []
    task_args = json.loads(task.args)
    for task_attr_name in task_args:
        # 布尔型只添加参数名
        if isinstance(task_args[task_attr_name], bool):
            if task_args[task_attr_name]:
                ops_args.append('%s' % str(task_attr_name))
        # 控制不添加
        elif not task_args[task_attr_name]:  # 如果参数值为空，则都不添加
            pass
        # json类型直接导入序列化以后的
        elif isinstance(task_args[task_attr_name], dict) or isinstance(
            task_args[task_attr_name], list
        ):
            ops_args.append('%s' % str(task_attr_name))
            args_values = json.dumps(task_args[task_attr_name], ensure_ascii=False)
            ops_args.append('%s' % args_values)

        # 其他的直接添加
        elif task_attr_name not in ['images', 'workdir']:
            ops_args.append('%s' % str(task_attr_name))
            args_values = task_args[task_attr_name]
            ops_args.append(
                '%s' % str(args_values)
            )  # 这里应该对不同类型的参数名称做不同的参数处理，比如bool型，只有参数，没有值

    # 设置环境变量
    container_envs = []
    if task.job_template.env:
        envs = re.split('\r|\n', task.job_template.env)
        envs = [env.strip() for env in envs if env.strip()]
        for env in envs:
            env_key, env_value = env.split('=')[0], env.split('=')[1]
            container_envs.append((env_key, env_value))

    # 设置全局环境变量
    # for global_env_key in global_envs:
    #     container_envs.append((global_env_key, global_envs[global_env_key]))

    # 设置task的默认环境变量
    container_envs.append(('KFJ_TASK_ID', str(task.id)))
    container_envs.append(('KFJ_TASK_NAME', str(task.name)))
    container_envs.append(('KFJ_TASK_NODE_SELECTOR', str(task.get_node_selector())))
    container_envs.append(('KFJ_TASK_VOLUME_MOUNT', str(task.volume_mount)))
    container_envs.append(('KFJ_TASK_IMAGES', str(task.job_template.images)))
    container_envs.append(('KFJ_TASK_RESOURCE_CPU', str(task.resource_cpu)))
    container_envs.append(('KFJ_TASK_RESOURCE_MEMORY', str(task.resource_memory)))
    container_envs.append(('KFJ_TASK_RESOURCE_GPU', str(task.resource_gpu.replace('+', ''))))
    container_envs.append(('USERNAME', pipeline.created_by.username))
    # if hubsecret_list:
    #     container_envs.append(("HUBSECRET", ','.join(hubsecret_list)))

    # 创建工作目录
    working_dir = None
    if task.job_template.workdir and task.job_template.workdir.strip():
        working_dir = task.job_template.workdir.strip()
    if task.working_dir and task.working_dir.strip():
        working_dir = task.working_dir.strip()

    # 配置启动命令
    task_command = ''

    if task.command:
        commands = re.split('\r|\n', task.command)
        commands = [command.strip() for command in commands if command.strip()]
        if task_command:
            task_command += ' && ' + ' && '.join(commands)
        else:
            task_command += ' && '.join(commands)

    job_template_entrypoint = (
        task.job_template.entrypoint.strip() if task.job_template.entrypoint else ''
    )

    command = None
    if job_template_entrypoint:
        command = job_template_entrypoint

    if task_command:
        command = task_command

    images = task.job_template.images.name
    command = command.split(' ') if command else []
    command = [com for com in command if com]
    arguments = ops_args
    json.loads(task.outputs) if task.outputs and json.loads(task.outputs) else None

    # 如果模板配置了images参数，那直接用模板的这个参数
    if json.loads(task.args).get('images', ''):
        images = json.loads(task.args).get('images')

    # 自定义节点
    if task.job_template.name == conf.get('CUSTOMIZE_JOB'):
        working_dir = json.loads(task.args).get('workdir')
        command = ['bash', '-c', json.loads(task.args).get('command')]
        arguments = []

    device_res = DeviceResource.from_str(task.resource_gpu)
    # 添加用户自定义挂载
    k8s_volumes = []
    k8s_volume_mounts = []
    task.volume_mount = task.volume_mount.strip() if task.volume_mount else ''
    if task.volume_mount:
        try:
            k8s_volumes, k8s_volume_mounts = py_k8s.K8s.get_volume_mounts(
                task.volume_mount, pipeline.created_by.username, resource_gpu=device_res
            )
        except Exception as e:
            print(e)

    # 添加node selector
    node_selector = {}
    for selector in re.split(',|;|\n|\t', task.get_node_selector()):
        selector = selector.replace(' ', '')
        if '=' in selector:
            node_selector[selector.strip().split('=')[0].strip()] = (
                selector.strip().split('=')[1].strip()
            )
    if device_res:
        for k, v in device_res.node_selector.items():
            node_selector[k] = v

    # 添加pod label
    pod_label = {
        'pipeline-id': str(pipeline.id),
        'pipeline-name': str(pipeline.name),
        'task-id': str(task.id),
        'task-name': str(task.name),
        # "run-id": global_envs.get('KFJ_RUN_ID', ''),
        'run-rtx': (
            g.user.username if g and g.user and g.user.username else pipeline.created_by.username
        ),
        'pipeline-rtx': pipeline.created_by.username,
    }
    pod_annotations = {
        'project': pipeline.project.name,
        'pipeline': pipeline.describe,
        'task': task.label,
        'job-template': task.job_template.describe,
    }

    # 设置资源限制
    resource_cpu = (
        task.job_template.get_env('TASK_RESOURCE_CPU')
        if task.job_template.get_env('TASK_RESOURCE_CPU')
        else task.resource_cpu
    )
    resource_gpu = (
        task.job_template.get_env('TASK_RESOURCE_GPU')
        if task.job_template.get_env('TASK_RESOURCE_GPU')
        else task.resource_gpu
    )

    resource_memory = (
        task.job_template.get_env('TASK_RESOURCE_MEMORY')
        if task.job_template.get_env('TASK_RESOURCE_MEMORY')
        else task.resource_memory
    )

    resources_requests = resources_limits = {}

    if resource_memory:
        if '~' not in resource_memory:
            resources_requests['memory'] = resource_memory
            resources_limits['memory'] = resource_memory
        else:
            resources_requests['memory'] = resource_memory.split('~')[0]
            resources_limits['memory'] = resource_memory.split('~')[1]

    if resource_cpu:
        if '~' not in resource_cpu:
            resources_requests['cpu'] = resource_cpu
            resources_limits['cpu'] = resource_cpu

        else:
            resources_requests['cpu'] = resource_cpu.split('~')[0]
            resources_limits['cpu'] = resource_cpu.split('~')[1]

    if resource_gpu:
        if device_res:
            for k, v in device_res.resources_requests.items():
                resources_requests[k] = v
            for k, v in device_res.resources_limits.items():
                resources_limits[k] = v

    # 配置host
    hostAliases = {}

    # global_hostAliases = conf.get('HOSTALIASES', '')
    # # global_hostAliases = ''
    # if task_temp.job_template.hostAliases:
    #     global_hostAliases += "\n" + task_temp.job_template.hostAliases
    # if global_hostAliases:
    #     hostAliases_list = re.split('\r|\n', global_hostAliases)
    #     hostAliases_list = [host.strip() for host in hostAliases_list if host.strip()]
    #     for row in hostAliases_list:
    #         hosts = row.strip().split(' ')
    #         hosts = [host for host in hosts if host]
    #         if len(hosts) > 1:
    #             hostAliases[hosts[1]] = hosts[0]

    if task.skip:
        command = ['echo', 'skip']
        arguments = None
        resources_requests = None
        resources_limits = None

    task_template = {
        # 'schedulerName': 'volcano',
        'name': task.name,  # 因为同一个
        'outputs': {'artifacts': []},
        'container': {
            'name': task.name + '-' + uuid.uuid4().hex[:4],
            'ports': [],
            'command': command,
            'args': arguments,
            'env': [{'name': item[0], 'value': item[1]} for item in container_envs],
            'image': images,
            'resources': {'limits': resources_limits, 'requests': resources_requests},
            'volumeMounts': k8s_volume_mounts,
            'workingDir': working_dir,
            'imagePullPolicy': conf.get('IMAGE_PULL_POLICY', 'Always'),
        },
        'nodeSelector': node_selector,
        'securityContext': {'privileged': True if task.job_template.privileged else False},
        'affinity': {
            'podAntiAffinity': {
                'preferredDuringSchedulingIgnoredDuringExecution': [
                    {
                        'podAffinityTerm': {
                            'labelSelector': {'matchLabels': {'pipeline-id': str(pipeline.id)}},
                            'topologyKey': 'kubernetes.io/hostname',
                        },
                        'weight': 80,
                    }
                ]
            }
        },
        'metadata': {'labels': pod_label, 'annotations': pod_annotations},
        'retryStrategy': {'limit': int(task.retry)} if task.retry else None,
        'volumes': k8s_volumes,
        'hostAliases': [
            {'hostnames': [hostname], 'ip': hostAliases[hostname]} for hostname in hostAliases
        ],
        'activeDeadlineSeconds': task.timeout if task.timeout else None,
    }

    # 统一添加一些固定环境变量，比如hostip，podip等
    task_template['container']['env'].append(
        {
            'name': 'K8S_NODE_NAME',
            'valueFrom': {'fieldRef': {'apiVersion': 'v1', 'fieldPath': 'spec.nodeName'}},
        }
    )
    task_template['container']['env'].append(
        {
            'name': 'K8S_POD_IP',
            'valueFrom': {'fieldRef': {'apiVersion': 'v1', 'fieldPath': 'status.podIP'}},
        }
    )
    task_template['container']['env'].append(
        {
            'name': 'K8S_HOST_IP',
            'valueFrom': {'fieldRef': {'apiVersion': 'v1', 'fieldPath': 'status.hostIP'}},
        }
    )
    task_template['container']['env'].append(
        {
            'name': 'K8S_POD_NAME',
            'valueFrom': {'fieldRef': {'apiVersion': 'v1', 'fieldPath': 'metadata.name'}},
        }
    )

    return task_template
